#!/usr/bin/env python3
"""
真正的 AI 驱动 MCP 客户端
集成大语言模型进行智能决策
"""

import asyncio
import aiohttp
import json
import logging
from typing import Dict, Any, Optional, List
from datetime import datetime
import openai
import os

class AIPlaywrightMCPClient:
    """AI 驱动的 Playwright MCP 客户端"""
    
    def __init__(self, mcp_server_url: str = "http://localhost:3000", openai_api_key: str = None):
        self.mcp_server_url = mcp_server_url
        self.session = None
        self.logger = logging.getLogger(__name__)
        
        # 初始化 OpenAI 客户端
        self.openai_client = openai.AsyncOpenAI(
            api_key=openai_api_key or os.getenv('OPENAI_API_KEY')
        )
        
        # 对话历史，用于保持上下文
        self.conversation_history = []
        
    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()
    
    async def get_ai_decision(self, context: str, options: List[str]) -> str:
        """使用 AI 做决策"""
        try:
            messages = [
                {
                    "role": "system", 
                    "content": """你是一个专业的网页自动化助手。你需要根据当前页面状态和用户目标，
                    选择最合适的操作。请始终选择最可能成功的选项，并简要说明原因。"""
                },
                {
                    "role": "user",
                    "content": f"""
                    当前情况: {context}
                    
                    可选操作:
                    {chr(10).join([f"{i+1}. {opt}" for i, opt in enumerate(options)])}
                    
                    请选择最佳操作（只返回数字）并简要说明原因。
                    格式: 数字|原因
                    """
                }
            ]
            
            # 添加历史对话以保持上下文
            if self.conversation_history:
                messages.extend(self.conversation_history[-4:])  # 保持最近4轮对话
            
            response = await self.openai_client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                temperature=0.3,
                max_tokens=200
            )
            
            decision = response.choices[0].message.content.strip()
            
            # 记录对话历史
            self.conversation_history.extend([
                {"role": "user", "content": f"情况: {context}"},
                {"role": "assistant", "content": decision}
            ])
            
            self.logger.info(f"🤖 AI 决策: {decision}")
            return decision
            
        except Exception as e:
            self.logger.error(f"❌ AI 决策失败: {e}")
            return "1|默认选择第一个选项"
    
    async def smart_element_find(self, description: str, page_snapshot: str) -> Dict[str, Any]:
        """AI 智能元素识别"""
        try:
            messages = [
                {
                    "role": "system",
                    "content": """你是一个网页元素识别专家。根据页面快照和元素描述，
                    找出最合适的选择器。返回 JSON 格式：
                    {"selector": "最佳选择器", "confidence": 0.9, "reason": "选择原因"}"""
                },
                {
                    "role": "user", 
                    "content": f"""
                    页面快照:
                    {page_snapshot[:2000]}...
                    
                    要找的元素: {description}
                    
                    请分析并返回最佳选择器。
                    """
                }
            ]
            
            response = await self.openai_client.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                temperature=0.1,
                max_tokens=300
            )
            
            result = response.choices[0].message.content.strip()
            
            # 尝试解析 JSON
            try:
                return json.loads(result)
            except:
                # 如果不是 JSON，返回默认结果
                return {
                    "selector": description,
                    "confidence": 0.5,
                    "reason": "AI 解析失败，使用原始描述"
                }
                
        except Exception as e:
            self.logger.error(f"❌ 智能元素识别失败: {e}")
            return {
                "selector": description,
                "confidence": 0.3,
                "reason": f"识别失败: {e}"
            }
    
    async def send_mcp_tool_call(self, tool_name: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """发送 MCP 工具调用"""
        try:
            # 构造标准 MCP 消息
            mcp_message = {
                "jsonrpc": "2.0",
                "id": f"req_{int(datetime.now().timestamp() * 1000)}",
                "method": "tools/call",
                "params": {
                    "name": tool_name,
                    "arguments": parameters
                }
            }
            
            self.logger.info(f"🔧 调用 MCP 工具: {tool_name}")
            self.logger.debug(f"参数: {parameters}")
            
            # 发送到 MCP 服务器
            async with self.session.post(
                f"{self.mcp_server_url}/messages",
                json=mcp_message,
                headers={
                    "Content-Type": "application/json",
                    "Accept": "text/event-stream"
                }
            ) as resp:
                
                if resp.status == 200:
                    # 处理 SSE 响应
                    result = await self._handle_sse_response(resp)
                    self.logger.info(f"✅ MCP 工具调用成功: {tool_name}")
                    return result
                else:
                    error_text = await resp.text()
                    self.logger.error(f"❌ MCP 工具调用失败: {resp.status} - {error_text}")
                    return {"error": f"HTTP {resp.status}: {error_text}"}
                    
        except Exception as e:
            self.logger.error(f"❌ MCP 工具调用异常: {e}")
            return {"error": str(e)}
    
    async def _handle_sse_response(self, response) -> Dict[str, Any]:
        """处理 SSE 响应"""
        try:
            content = await response.text()
            
            # 简单的 SSE 解析（实际应该更完善）
            lines = content.split('\n')
            for line in lines:
                if line.startswith('data: '):
                    data = line[6:]  # 移除 'data: ' 前缀
                    if data.strip():
                        return json.loads(data)
            
            return {"success": True, "data": "操作完成"}
            
        except Exception as e:
            self.logger.error(f"❌ SSE 响应解析失败: {e}")
            return {"error": f"响应解析失败: {e}"}
    
    async def intelligent_navigate(self, target_description: str) -> bool:
        """智能导航"""
        # 1. 获取当前页面状态
        snapshot_result = await self.send_mcp_tool_call("mcp_Playwright_browser_snapshot", {
            "random_string": "nav_check"
        })
        
        if "error" in snapshot_result:
            self.logger.error(f"获取页面快照失败: {snapshot_result['error']}")
            return False
        
        current_state = snapshot_result.get("data", "")
        
        # 2. AI 分析当前状态和目标
        context = f"当前页面状态: {current_state[:500]}...\n目标: {target_description}"
        options = [
            "直接导航到目标URL",
            "在当前页面查找导航链接",
            "使用搜索功能",
            "返回首页重新开始"
        ]
        
        decision = await self.get_ai_decision(context, options)
        choice_num = int(decision.split('|')[0]) if '|' in decision else 1
        
        # 3. 根据 AI 决策执行操作
        if choice_num == 1:
            # 直接导航
            if "小红书" in target_description:
                result = await self.send_mcp_tool_call("mcp_Playwright_browser_navigate", {
                    "url": "https://creator.xiaohongshu.com"
                })
                return "error" not in result
        
        # 其他选择的实现...
        return True
    
    async def intelligent_click(self, element_description: str) -> bool:
        """智能点击"""
        # 1. 获取页面快照
        snapshot_result = await self.send_mcp_tool_call("mcp_Playwright_browser_snapshot", {
            "random_string": "click_analysis"
        })
        
        if "error" in snapshot_result:
            return False
        
        page_snapshot = snapshot_result.get("data", "")
        
        # 2. AI 智能识别元素
        element_info = await self.smart_element_find(element_description, page_snapshot)
        
        if element_info["confidence"] < 0.6:
            self.logger.warning(f"⚠️ 元素识别置信度较低: {element_info['confidence']}")
        
        # 3. 执行点击
        result = await self.send_mcp_tool_call("mcp_Playwright_browser_click", {
            "element": element_description,
            "ref": "ai-detected"
        })
        
        return "error" not in result
    
    async def intelligent_input(self, field_description: str, content: str) -> bool:
        """智能输入"""
        # 1. 获取页面状态
        snapshot_result = await self.send_mcp_tool_call("mcp_Playwright_browser_snapshot", {
            "random_string": "input_analysis"
        })
        
        if "error" in snapshot_result:
            return False
        
        # 2. AI 识别输入字段
        page_snapshot = snapshot_result.get("data", "")
        element_info = await self.smart_element_find(field_description, page_snapshot)
        
        # 3. 执行输入
        result = await self.send_mcp_tool_call("mcp_Playwright_browser_type", {
            "element": field_description,
            "ref": "ai-detected",
            "text": content
        })
        
        return "error" not in result

class IntelligentXiaohongshuPublisher:
    """智能小红书发布器"""
    
    def __init__(self, ai_client: AIPlaywrightMCPClient):
        self.ai_client = ai_client
        self.logger = logging.getLogger(__name__)
    
    async def publish_with_ai_guidance(self, content_data: Dict[str, str], image_path: str = None):
        """AI 指导的发布流程"""
        try:
            self.logger.info("🤖 开始 AI 指导的发布流程...")
            
            # 1. 智能导航到发布页面
            if not await self.ai_client.intelligent_navigate("小红书图文发布页面"):
                raise Exception("智能导航失败")
            
            # 2. AI 等待页面加载完成
            await asyncio.sleep(3)
            
            # 3. 智能选择发布类型
            if not await self.ai_client.intelligent_click("图文发布按钮"):
                self.logger.warning("⚠️ 未找到图文发布按钮，尝试其他方式")
                if not await self.ai_client.intelligent_click("上传图文"):
                    raise Exception("无法找到发布入口")
            
            # 4. 智能上传图片
            if image_path:
                result = await self.ai_client.send_mcp_tool_call("mcp_Playwright_browser_file_upload", {
                    "paths": [image_path]
                })
                if "error" in result:
                    self.logger.warning(f"图片上传失败: {result['error']}")
                else:
                    await asyncio.sleep(3)  # 等待上传完成
            
            # 5. 智能填写标题
            if not await self.ai_client.intelligent_input("标题输入框", content_data["title"]):
                self.logger.warning("⚠️ 标题输入失败")
            
            # 6. 智能填写正文
            if not await self.ai_client.intelligent_input("正文输入框", content_data["content"]):
                self.logger.warning("⚠️ 正文输入失败")
            
            # 7. 智能发布
            if not await self.ai_client.intelligent_click("发布按钮"):
                raise Exception("发布失败")
            
            # 8. AI 验证发布结果
            await asyncio.sleep(5)
            
            # 获取最终状态
            final_snapshot = await self.ai_client.send_mcp_tool_call("mcp_Playwright_browser_snapshot", {
                "random_string": "final_check"
            })
            
            # AI 分析是否发布成功
            if "success" in str(final_snapshot).lower() or "发布成功" in str(final_snapshot):
                self.logger.info("🎉 AI 确认发布成功！")
                return True
            else:
                self.logger.warning("⚠️ 发布状态不明确")
                return False
                
        except Exception as e:
            self.logger.error(f"❌ AI 指导发布失败: {e}")
            return False

async def main():
    """主函数"""
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
    
    # 检查 API Key
    if not os.getenv('OPENAI_API_KEY'):
        print("❌ 请设置 OPENAI_API_KEY 环境变量")
        return
    
    async with AIPlaywrightMCPClient() as ai_client:
        publisher = IntelligentXiaohongshuPublisher(ai_client)
        
        content = {
            "title": "🤖 AI 驱动的自动化发布测试",
            "content": """这是一个真正由 AI 驱动的自动化发布测试！

🧠 特色功能：
• AI 智能识别页面元素
• 自适应页面变化
• 智能错误恢复
• 语义理解用户意图

#AI自动化 #智能发布 #技术测试"""
        }
        
        success = await publisher.publish_with_ai_guidance(
            content,
            "/Users/samciu/Sites/cgb-work/playwright-mcp-survey/python-playwright-xhs/sample/background.jpg"
        )
        
        if success:
            print("✅ AI 驱动的发布任务完成！")
        else:
            print("❌ 发布任务失败")

if __name__ == "__main__":
    asyncio.run(main())
